#coding:utf-8
import socket
import xlwt
import xlrd
import time
from multiprocessing import Pool

rootpath = "D:\\Go"
ipList = []
ipv6List = []
with open(rootpath + "/ip.txt", "r") as fip:
    for ip in fip.readlines():
        ip = ip.strip()
        ipList.append(ip)

with open(rootpath + "/ipv6.txt", "r") as fipv6:
    for ipv6 in fipv6.readlines():
        ipv6 = ipv6.strip()
        ipv6List.append(ipv6)       

def domain2ip(urlList):
    urlList[1] = urlList[1].strip()
    # urlList = url.split( )
    try:
        ip = socket.gethostbyname(str(urlList[1]))
        if ip in ipList:
            tips = "是"
        else:
            tips = "否"
    except:
        print(urlList[1] + " to IP ERROR ")
        ip = "未知"
        tips = "未知"

    try:
        result = socket.getaddrinfo(str(urlList[1]), None, socket.AF_INET6)
        ipv6 = result[0][4][0]
        if ipv6 in ipv6List:
            tipsv6 = "是"
        else:
            tipsv6 = "否"
    except:
        ipv6 = "unknown"
        tipsv6 = "未知"
    
    return [urlList[0], urlList[1], str(ip), str(tips), str(ipv6), str(tipsv6)]

def writeExcelFile(filename, content):
    workbook = xlwt.Workbook(encoding = 'utf-8')
    booksheet = workbook.add_sheet('域名解析表', cell_overwrite_ok = True)
    length = len(content)
    for i in range(0, length):
        for j in range(0, len(content[i])): # 如果长度确定，也可以直接写数字，比如 range(0, 3)
            booksheet.write(i, j, content[i][j])

    workbook.save(filename)

if __name__ == '__main__':
    # domains
    # allUrls = []
    # with open(rootpath + "/domain.txt", "r", encoding='utf-8') as urllist:
    #     allUrls = urllist.readlines()

    data = xlrd.open_workbook(rootpath + '/test.xls') # 不支持xlsx格式
    # sheet = data.sheet_by_name('Sheet1')
    sheet = data.sheet_by_index(0)

    datalist = []
    for row in range(1, sheet.nrows):
        rowdata = sheet.row_values(row)
        
        domains = rowdata[2].strip()
        domains = domains.split(';')

        for domain in domains:
            if domain != '':
                datalist.append([rowdata[0], domain])

    # resultList = []
    # for url in allUrls:
    #     resultList.append(domain2ip(url))
    p = Pool(8) # 建议设置成CPU核数
    resultList = p.map(domain2ip, datalist)
    p.close()
    p.join()

    resultList.insert(0, ['公司名', '域名', 'IP', 'IPv4是否在公司', 'IPv6', 'IPv6是否在公司'])
    filename = rootpath + '/域名解析表' + time.strftime("%Y-%m-%d-%H%M%S") + '.xls'
    writeExcelFile(filename, resultList)

    # print(resultList)
    # write the result to file
    # with open(rootpath + "/resolve.txt", "w") as resovelist:
    #     resovelist.writelines("\n" . join(resultList))

    print("Well Done! Resolve Completed!")